from ultralytics import YOLO
import os
import json  # 导入json模块用于保存验证结果
import shutil  # 导入shutil模块用于删除目录
import matplotlib.pyplot as plt  # 导入matplotlib用于绘制柱形图
import datetime  # 导入datetime模块用于处理时间戳
def validate_model(model_path,t=0):
    """
    验证模型并把验证结果保存为json在模型目录内，
    如果目录内已经有验证结果，则不要验证，直接读取验证结果并打印
    """
    # 获取模型目录和JSON文件路径
    model_dir = os.path.dirname(model_path)
    json_path = os.path.join(model_dir, "validation_results.json")
    
    # 检查JSON文件是否已存在
    if os.path.exists(json_path):
        try:
            # 读取已存在的验证结果
            with open(json_path, 'r', encoding='utf-8') as f:
                validation_results = json.load(f)
            
            print(f"mAP50: {validation_results['mAP50']}")
            print(f"mAP50-95: {validation_results['mAP50-95']}")
            print(f"精确率: {validation_results['precision']}")
            print(f"召回率: {validation_results['recall']}")
            print(f"时间消耗: {validation_results['time-consumption']}")
            return  # 直接返回，不重新验证
        except Exception as e:
            print(f"读取验证结果时出错: {str(e)}")
            
    
    # 如果JSON文件不存在或读取失败，则进行模型验证
    print('首次加载模型,正在验证...')
    model=YOLO(model_path)
    metrics = model.val(
        name='ttm_val',
        
    )  # no arguments needed, dataset and settings remembered
    print(f"mAP50: {metrics.box.map50:.4f}")  # map50
    print(f"mAP50-95: {metrics.box.map:.4f}")  # map50-95
    print(f"精确率: {metrics.box.p[0]:.4f}")  # precision
    print(f"召回率: {metrics.box.r[0]:.4f}")  # recall
    print(f"时间消耗: {t}分")  # time consumption
    
    # 创建包含验证结果的字典
    validation_results = {
        "mAP50": round(float(metrics.box.map50), 4),
        "mAP50-95": round(float(metrics.box.map), 4),
        "precision": round(float(metrics.box.p[0]), 4),
        "recall": round(float(metrics.box.r[0]), 4),
        'time-consumption': t,
    }
    
    # 保存到模型目录
    try:
        with open(json_path, 'w', encoding='utf-8') as f:
            json.dump(validation_results, f, ensure_ascii=False, indent=2)
        print(f"验证结果已保存到: {json_path}")
    except Exception as e:
        print(f"保存验证结果时出错: {str(e)}")
    
    # 删除验证结果目录
    try:
        if os.path.exists('E:/Py/visual/runs/detect/ttm_val'):
            shutil.rmtree('E:/Py/visual/runs/detect/ttm_val')
            # print("验证结果目录已成功删除")
    except Exception as e:
        print(f"删除目录时出错: {str(e)}")

    return validation_results


def show_all_val(directory,block=True):
    """
    读取目录内所有模型的四项性能指标，如果没有验证结果则调用validate_model进行验证，
    然后使用matplotlib绘制柱形图进行比较
    """
    # 只收集名为best.pt的模型文件
    model_files = []
    for root, dirs, files in os.walk(directory):
        for file in files:
            if file == 'best.pt':  # 只收集文件名等于best.pt的文件
                model_files.append(os.path.join(root, file))
    
    if not model_files:
        print(f"在目录 {directory} 中未找到.pt模型文件")
        return
    
    # 辅助函数：从模型路径中提取时间戳
    def get_model_timestamp(model_path):
        model_dir = os.path.dirname(model_path)
        parent_dir = os.path.basename(os.path.dirname(model_dir))
        parts = parent_dir.split('_')
        if len(parts) < 2:
            return datetime.datetime.min
        try:
            # 提取时间戳部分（假设格式为YYYYMMDD_HHMMSS）
            date_part = parts[-2]
            time_part = parts[-1]
            timestamp_str = f"{date_part}_{time_part}"
            return datetime.datetime.strptime(timestamp_str, "%Y%m%d_%H%M%S")
        except (ValueError, IndexError):
            # 如果解析失败，返回一个非常早的日期
            return datetime.datetime.min
    
    # 按时间戳对模型文件进行排序，最新的模型排在最后（将显示在最右侧）
    model_files.sort(key=get_model_timestamp)
    
    # 收集每个模型的验证结果
    model_results = []
    model_names = []
    
    # 辅助函数：从模型名称中截去时间戳部分
    def remove_timestamp_from_name(model_name):
        parts = model_name.split('_')
        # 检查最后两个部分是否是时间戳格式（YYYYMMDD_HHMMSS）
        if len(parts) >= 2:
            # 检查倒数第二部分是否是8位数字（日期）
            if len(parts[-2]) == 8 and parts[-2].isdigit():
                # 检查最后一部分是否是6位数字（时间）
                if len(parts[-1]) == 6 and parts[-1].isdigit():
                    # 移除时间戳部分
                    return '_'.join(parts[:-2])
        # 如果不符合时间戳格式，返回原始名称
        return model_name
    
    for idx, model_path in enumerate(model_files):
        # 获取模型所在的父目录名称（而不是模型文件名）
        model_dir = os.path.dirname(model_path)
        parent_dir = os.path.basename(os.path.dirname(model_dir))  # 获取权重目录的父目录
        
        # 截去时间戳部分以缩减名称长度
        shortened_name = remove_timestamp_from_name(parent_dir)
        
        # 生成显示名称
        model_display_name = f"{idx}: {shortened_name}"
        model_names.append(model_display_name)
        
        print(f"处理模型 {model_display_name} (路径: {model_path})")
        
        # 获取模型目录和JSON文件路径
        model_dir = os.path.dirname(model_path)
        json_path = os.path.join(model_dir, "validation_results.json")
        
        # 如果JSON文件不存在，则调用validate_model进行验证
        if not os.path.exists(json_path):
            # validate_model函数会保存验证结果
            validate_model(model_path)
        
        # 读取验证结果
        try:
            with open(json_path, 'r', encoding='utf-8') as f:
                validation_results = json.load(f)
            model_results.append(validation_results)
        except Exception as e:
            print(f"读取模型 {model_display_name} 的验证结果时出错: {str(e)}")
            # 添加默认值，避免图表绘制失败
            model_results.append({
                "mAP50": 0,
                "mAP50-95": 0,
                "precision": 0,
                "recall": 0,
                
            })
    
    # 提取性能指标数据
    metrics = ["mAP50", "mAP50-95", "precision", "recall"]
    data = {metric: [result[metric] for result in model_results] for metric in metrics}
    
    # 找出mAP50-95最高的模型索引
    max_map_idx = data["mAP50-95"].index(max(data["mAP50-95"]))
    
    # 创建柱形图
    fig, ax = plt.subplots(figsize=(9, 6))
    
    # 设置条形宽度
    width = 0.1
    
    # 设置x轴位置
    x = [i + width * j for j in range(len(model_names)) for i in range(len(metrics))]
    
    # 创建一个双Y轴的图表，左侧用于柱状图，右侧用于折线图
    ax2 = ax.twinx()
    
    # 绘制每个指标的条形图和折线图
    colors = ['blue', 'green', 'red', 'purple']
    # markers = ['o', 's', '^', 'D']  # 不同的标记样式用于折线图
    bars = []
    lines = []
    
    # 亮绿色设置
    highlight_color = '#00FF00'  # 亮绿色
    
    for i, (metric, values) in enumerate(data.items()):
        x_pos = [j + i * width for j in range(len(model_names))]
        
        # 使用普通颜色绘制柱状图，不修改最佳模型柱子的颜色
        bar = ax.bar(x_pos, values, width=width, label=metric, color=colors[i % len(colors)], alpha=0.7)
        bars.append(bar)
        
        # 确保折线图端点精确位于柱状图中心线上
        # 使用与柱状图相同的x坐标计算方式，但确保每个点都在柱中心
        # 核心优化：确保line_x_pos与柱状图的x_pos + width/2精确对应
        line_x_pos = [j + i * width + width/2 for j in range(len(model_names))]
        
        # 绘制折线图，增加线连接的平滑度，确保端点与柱中心对齐
        line, = ax2.plot(line_x_pos, values, color=colors[i % len(colors)], 
                        markersize=8, linewidth=2, label=f'{metric}_趋势')
        
        # 确保折线的每个点都明确显示在柱中心（通过调整标记大小增强可见性）
        lines.append(line)
        
        # 为每个柱子添加数值标签，所有标签使用黑色
        for j, rect in enumerate(bar):
            height = rect.get_height()
            ax.annotate(f'{height:.4f}',
                        xy=(rect.get_x() + rect.get_width() / 2, height),
                        xytext=(0, 3),  # 3点垂直偏移
                        textcoords="offset points",
                        ha='center', va='bottom',
                        fontsize=8,
                        rotation=45,
                        color='black')  # 所有标签使用黑色文本
    
    # 在最佳模型上方添加绿色的"BEST!"文本
    # 计算最佳模型的最高柱子高度作为文本位置的基准
    max_height = 0
    for i in range(len(metrics)):
        height = bars[i][max_map_idx].get_height()
        if height > max_height:
            max_height = height
    
    # 计算最佳模型柱子组的中心点位置
    center_pos = max_map_idx + (len(metrics) - 1) * width / 2
    
    # 添加BEST!文本
    ax.text(
        center_pos,
        max_height * 1.05,  # 略高于最高柱子
        'BEST!',
        ha='center',
        va='bottom',
        color=highlight_color,
        fontsize=12,
        fontweight='bold'
    )
    
    # 合并图例（柱状图和折线图）
    combined_labels = [metric for metric in metrics] 
    combined_handles = [bar[0] for bar in bars] + lines
    ax.legend(combined_handles, combined_labels, loc='upper left')
    
    # 设置右侧Y轴标签
    ax2.set_ylabel('Trend Value', fontsize=10)
    # 确保右侧Y轴范围与左侧一致
    ax2.set_ylim(ax.get_ylim())
    
    # 设置x轴标签，为mAP50-95最高的模型使用亮绿色标签
    xticks = [j + width * (len(metrics) - 1) / 2 for j in range(len(model_names))]
    ax.set_xticks(xticks)
    
    # 设置标签颜色 - 修复设置方法，先设置标签再修改颜色
    ax.set_xticklabels(model_names, rotation=45, ha='right')
    
    # 设置标签颜色（所有标签使用黑色）
    for label in ax.get_xticklabels():
        label.set_color('black')  # 所有标签使用黑色
    
    # 设置图表标题和标签
    ax.set_title('model_traits_comparison')
    ax.set_ylabel('metric_value')

    # 调整布局
    plt.tight_layout()
    
    # 保存图表到当前目录 - 使用SVG矢量格式，放大后不会失真
    # 其他可选的矢量格式：'pdf', 'eps', 'ps'
    plt.savefig('model_performance_comparison.svg', format='svg')
    print("模型性能比较图表已保存为 model_performance_comparison.svg (矢量格式)")
        
    # 设置键盘事件处理函数
    def on_key_press(event):
        # 支持ESC键(27)、q键(113)和空格键(32)关闭图形
        if event.key == 'escape' or event.key == 'q' or event.key == ' ':
            
            plt.close()
    
    # 连接键盘事件
    fig.canvas.mpl_connect('key_press_event', on_key_press)
    print("按ESC键、Q键或空格键关闭图表。")
    if block:
        # 显示图表
        plt.show()
    else:
        # 显示2秒后关闭图像
        plt.show(block=False)
        plt.pause(2)
        plt.close()
        




if __name__ == '__main__':
    
    show_all_val('E:/Py/visual/runs/detect',block=True)
